import requests
import re
import time
import json
import os
import brotli
import argparse
from pprint import pprint
from urllib.parse import unquote
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from seleniumwire import webdriver
from io import BytesIO

# 创建 ChromeOptions 对象
options = Options()
options.add_argument('--headless')  # 可选：无头模式
options.add_argument('--disable-gpu')  # 可选：禁用 GPU 加速
options.add_argument('user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36')

# 初始化 WebDriver（需要安装相应的浏览器驱动）
service = Service('C:/Users/JIA/AppData/Local/Google/Chrome/Application/chromedriver.exe')
driver = webdriver.Chrome(service=service, options=options)
headers1 = {
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36'
}

true_url = "https://www.douyin.com/video/"

# 获取视频链接
def get_dy_url(text):
    dy_link = filter_url(text)
    stats1 = requests.get(url=dy_link, headers=headers1)
    locationlink = stats1.url
    pp1 = requests.get(url=locationlink, headers=headers1)
    referer_url = pp1.url
    videonum = re.search(r"\d+", referer_url).group()
    tureUrl = true_url + videonum
    # print(tureUrl)
    return parse(tureUrl)

# 正则提取视频链接
def filter_url(text):
    regex = r'https?://[\w-]+(\.[\w-]+)+(/[\w.-]*(\?\w+=([\w%-]+)(&\w+=([\w%-]+)*)?)?)?'
    match = re.search(regex, text)
    if match:
        return match.group(0)
    return ""

# 爬取视频信息
def parse(tureUrl):
    driver.get(url=tureUrl)
    time.sleep(2)  # 根据需要调整等待时间
    pattern = r'https?://www\.douyin\.com/aweme/v1/web/aweme/detail.*'
    matched_responses = []
    for request in driver.requests:
        if re.match(pattern, request.url):
            matched_responses.append((request.url, request.response))

    if matched_responses:
        # 获取第一个匹配项
        url, response = matched_responses[0]
        if response:
            response_body = response.body
            if response.headers.get('Content-Encoding') == 'br':
                response_body = brotli.decompress(response_body).decode('utf-8')
            else:
                response_body = response_body.decode('utf-8')

            try:
                json_data = json.loads(response_body)
                bit_rate = json_data['aweme_detail']['video']['bit_rate'][0]
                video_url = bit_rate['play_addr']['url_list']
                desc = json_data['aweme_detail']['desc']
                driver.quit()  # 在此处关闭浏览器
                return download_video(video_url, 'D:/000-total_folder/200-project_refrence/project/local/Python/video', desc)

            except (KeyError, json.JSONDecodeError) as e:
                print("解析 JSON 时出错:", e)

    driver.quit()
    return None

# 下载视频
def download_video(video_url, download_folder, desc):
    name = remove_hash_and_after(desc)
    if not os.path.exists(download_folder):
        os.mkdir(download_folder)

    for i in range(len(video_url)):
        video_path = os.path.join(download_folder, f"{name}_{i}.mp4")
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36',
            'Referer': 'https://www.douyin.com/'
        }

        try:
            video_data = requests.get(url=video_url[i], headers=headers).content
            if len(video_data) > 0:
                with open(video_path, 'wb') as f:
                    f.write(video_data)
                # print(f"成功下载: {video_path}")
                return video_path  # 返回成功下载的视频路径

        except requests.exceptions.RequestException as e:
            print(f"下载失败: {video_url[i]}, 错误: {e}")

    return None  # 如果没有成功下载任何视频，返回 None


# 去掉#
def remove_hash_and_after(s):
    return s.split('#')[0]

# 主函数
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='解析抖音视频链接')
    parser.add_argument('inputurl', type=str, help='抖音视频链接')
    args = parser.parse_args()
    inputurl = args.inputurl
    t1 = time.time()
    downloaded_paths = get_dy_url(inputurl)  # 获取下载路径
    t2 = time.time()
    tt = t2 - t1
    # print('解析耗时(sec):', round(tt, 2))
    print('下载的文件路径:', downloaded_paths)  # 打印下载路径
